package com.study.chapter09;

import com.study.chapter05.source.ClickSource;
import com.study.entity.Event;
import org.apache.flink.api.common.eventtime.WatermarkStrategy;
import org.apache.flink.api.common.functions.AggregateFunction;
import org.apache.flink.api.common.functions.FlatMapFunction;
import org.apache.flink.api.common.functions.RichFlatMapFunction;
import org.apache.flink.api.common.state.AggregatingState;
import org.apache.flink.api.common.state.AggregatingStateDescriptor;
import org.apache.flink.api.common.state.ValueState;
import org.apache.flink.api.common.state.ValueStateDescriptor;
import org.apache.flink.api.common.typeinfo.Types;
import org.apache.flink.api.java.tuple.Tuple2;
import org.apache.flink.configuration.Configuration;
import org.apache.flink.streaming.api.datastream.SingleOutputStreamOperator;
import org.apache.flink.streaming.api.environment.StreamExecutionEnvironment;
import org.apache.flink.util.Collector;

import java.sql.Timestamp;
import java.time.Duration;

/**
 * @Description:
 * @Author: LiuQun
 * @Date: 2022/8/16 21:02
 */
public class AggregateStateTest {
    public static void main(String[] args) throws Exception {
        StreamExecutionEnvironment env = StreamExecutionEnvironment.getExecutionEnvironment();
        env.setParallelism(1);

        SingleOutputStreamOperator<Event> stream = env.addSource(new ClickSource())
                .assignTimestampsAndWatermarks(
                        WatermarkStrategy.<Event>forBoundedOutOfOrderness(Duration.ZERO)
                                .withTimestampAssigner((event, l) -> event.timestamp)
                );
        stream.print("input=>");

        //统计每个用户每5次访问的平均时间
        stream.keyBy(event -> event.user)
                .flatMap(new AggregateResult())
                .print();

        env.execute();
    }

    public static class AggregateResult extends RichFlatMapFunction<Event,String>{
        //定义一个聚合的状态，用来保存平均时间戳
        private AggregatingState<Event,Long> timeAggState;
        //定义一个值状态，用来保存用户访问的次数
        private ValueState<Long> countState;

        @Override
        public void open(Configuration parameters) throws Exception {
           timeAggState = getRuntimeContext().getAggregatingState(
                   //Tuple2中的f0存放时间戳之和，f1计数
                   new AggregatingStateDescriptor<Event, Tuple2<Long,Long>, Long>(
                           "time-agg",
                           new AggregateFunction<Event, Tuple2<Long, Long>, Long>() {
                               @Override
                               public Tuple2<Long, Long> createAccumulator() {
                                   return Tuple2.of(0L,0L);
                               }

                               @Override
                               public Tuple2<Long, Long> add(Event value, Tuple2<Long, Long> accumulator) {
                                   return Tuple2.of(value.timestamp + accumulator.f0,accumulator.f1+1);
                               }

                               @Override
                               public Long getResult(Tuple2<Long, Long> accumulator) {
                                   //输出平均时间戳
                                   return accumulator.f0 / accumulator.f1;
                               }

                               @Override
                               public Tuple2<Long, Long> merge(Tuple2<Long, Long> a, Tuple2<Long, Long> b) {
                                   return null;
                               }
                           },
                           Types.TUPLE(Types.LONG,Types.LONG)
                   )
           );
           countState = getRuntimeContext().getState(new ValueStateDescriptor<Long>("count",Long.class));

        }

        @Override
        public void flatMap(Event value, Collector<String> out) throws Exception {

            Long count = countState.value();
            if (count == null){
                count = 1L;
            }else {
                count++;
            }
            countState.update(count);
            timeAggState.add(value);
            //达到5次就输出结果，并清空状态
            if (count == 5){
                out.collect(value.user + " 平均时间戳：" + new Timestamp(timeAggState.get()));
                countState.clear();
            }
        }
    }
}
